{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SUknhuHKyc-E"
   },
   "source": [
    "# <center>OpenAI agent pattern: parallelization</center>\n",
    "\n",
    "A starter guide for building an agent with parallel LLM and tool calling using the `openai-agents` library.\n",
    "\n",
    "In this example, we are building a stock research agent. If we want to research 5 stocks, we can force the agent to run multiple tool calls, instead of sequentially."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n69HR7eJswNt"
   },
   "source": [
    "### Install Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install base libraries for OpenAI\n",
    "!pip install -q openai openai-agents yfinance\n",
    "\n",
    "# Install optional libraries for OpenInference/OpenTelemetry tracing\n",
    "!pip install -q arize-phoenix-otel openinference-instrumentation-openai-agents openinference-instrumentation-openai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jQnyEnJisyn3"
   },
   "source": [
    "### Setup Keys\n",
    "\n",
    "Add your OpenAI API key to the environment variable `OPENAI_API_KEY`.\n",
    "\n",
    "Copy your Phoenix `API_KEY` from your settings page at [app.phoenix.arize.com](https://app.phoenix.arize.com)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    os.environ[\"OPENAI_API_KEY\"] = getpass(\"🔑 Enter your OpenAI API key: \")\n",
    "\n",
    "if \"PHOENIX_API_KEY\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_API_KEY\"] = getpass(\"🔑 Enter your Phoenix API key: \")\n",
    "\n",
    "if \"PHOENIX_COLLECTOR_ENDPOINT\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_COLLECTOR_ENDPOINT\"] = getpass(\"🔑 Enter your Phoenix Collector Endpoint\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kfid5cE99yN5"
   },
   "source": [
    "### Setup Tracing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from phoenix.otel import register\n",
    "\n",
    "tracer_provider = register(\n",
    "    project_name=\"openai-agents\",\n",
    "    auto_instrument=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating your agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "import yfinance as yf\n",
    "\n",
    "\n",
    "def get_stock_data(ticker_symbol):\n",
    "    stock = yf.Ticker(ticker_symbol)\n",
    "    data = stock.info\n",
    "    return {\n",
    "        \"symbol\": data.get(\"symbol\"),\n",
    "        \"current_price\": data.get(\"currentPrice\"),\n",
    "        \"market_cap\": data.get(\"marketCap\"),\n",
    "        \"sector\": data.get(\"sector\"),\n",
    "        \"industry\": data.get(\"industry\"),\n",
    "        \"description\": data.get(\"longBusinessSummary\"),\n",
    "        \"trailing_pe\": data.get(\"trailingPE\"),\n",
    "        \"forward_pe\": data.get(\"forwardPE\"),\n",
    "        \"dividend_yield\": data.get(\"dividendYield\"),\n",
    "        \"beta\": data.get(\"beta\"),\n",
    "        \"fifty_two_week_high\": data.get(\"fiftyTwoWeekHigh\"),\n",
    "        \"fifty_two_week_low\": data.get(\"fiftyTwoWeekLow\"),\n",
    "        \"fifty_day_moving_average\": data.get(\"fiftyDayAverage\"),\n",
    "        \"two_hundred_day_moving_average\": data.get(\"twoHundredDayAverage\"),\n",
    "        \"recommendation_key\": data.get(\"recommendationKey\"),\n",
    "        \"revenue_growth\": data.get(\"revenueGrowth\"),\n",
    "        \"earnings_growth\": data.get(\"earningsGrowth\"),\n",
    "        \"profit_margins\": data.get(\"profitMargins\"),\n",
    "    }\n",
    "\n",
    "\n",
    "pprint(get_stock_data(\"AAPL\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "from textwrap import dedent\n",
    "\n",
    "from agents import Agent, Runner, TResponseInputItem, WebSearchTool, function_tool\n",
    "from agents.model_settings import ModelSettings\n",
    "\n",
    "\n",
    "@function_tool\n",
    "def get_stock_data_tool(ticker_symbol: str) -> dict:\n",
    "    \"\"\"\n",
    "    Get stock data for a given ticker symbol.\n",
    "    Args:\n",
    "        ticker_symbol: The ticker symbol of the stock to get data for.\n",
    "    Returns:\n",
    "        A dictionary containing stock data such as price, market cap, and more.\n",
    "    \"\"\"\n",
    "    return get_stock_data(ticker_symbol)\n",
    "\n",
    "\n",
    "research_agent = Agent(\n",
    "    name=\"FinancialSearchAgent\",\n",
    "    instructions=dedent(\n",
    "        \"\"\"You are a research assistant specializing in financial topics. Given a stock ticker, use web search to retrieve up‑to‑date context and produce a short summary of at most 50 words. Focus on key numbers, events, or quotes that will be useful to a financial analyst.\"\"\"\n",
    "    ),\n",
    "    model=\"gpt-4.1\",\n",
    "    tools=[WebSearchTool(), get_stock_data_tool],\n",
    "    model_settings=ModelSettings(tool_choice=\"required\", parallel_tool_calls=True),\n",
    ")\n",
    "\n",
    "user_input = input(\"Enter the stock tickers you want to research: \")\n",
    "input_items: list[TResponseInputItem] = [{\"content\": user_input, \"role\": \"user\"}]\n",
    "\n",
    "orchestrator = await Runner.run(research_agent, input_items)\n",
    "orchestrator_output = orchestrator.final_output\n",
    "pprint(orchestrator_output)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
